import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import os

def generate_image(prompt, negative_prompt="", num_inference_steps=50, guidance_scale=7.5):
    """使用 Stable Diffusion 生成图片"""
    # 检查本地模型是否存在
    model_path = "./stable-diffusion-v1-5"
    if not os.path.exists(model_path):
        print("本地模型不存在，请先运行 download_sd_model.py 下载模型")
        return None
    
    # 初始化模型
    pipe = StableDiffusionPipeline.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        use_safetensors=True
    )
    
    # 如果有 CUDA 则使用 GPU
    if torch.cuda.is_available():
        pipe = pipe.to("cuda")
    
    # 生成图片
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).images[0]
    
    return image

def main():
    # 示例使用
    prompt = "一只可爱的猫咪，坐在窗台上看着窗外的雨"
    negative_prompt = "模糊的, 扭曲的, 低质量的"
    
    try:
        print("正在生成图片，请稍候...")
        image = generate_image(prompt, negative_prompt)
        
        if image is not None:
            # 保存图片
            output_path = "generated_image.png"
            image.save(output_path)
            print(f"图片已生成并保存为: {output_path}")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 